#!/usr/bin/env python3
"""
将表格图片信息导入到QA_knowledge数据库的table_images表中
"""

import os
import sys
import json
import base64
import configparser
from pathlib import Path
from typing import List, Dict
import psycopg2
from psycopg2.extras import execute_values

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from table.config import DATASET_DIR, DATABASES


def get_qa_database_connection():
    """连接到QA_knowledge数据库"""
    config = configparser.ConfigParser()
    
    # 优先从table目录读取配置文件
    config_path = Path(__file__).parent / 'db_config.ini'
    
    if not config_path.exists():
        # 尝试从D_db/core目录读取
        config_path = project_root / 'D_db' / 'core' / 'db_config.ini'
    
    if not config_path.exists():
        raise FileNotFoundError(f"找不到配置文件: {config_path}")
    
    print(f"📄 使用配置文件: {config_path}")
    config.read(str(config_path), encoding='utf-8')
    
    try:
        conn = psycopg2.connect(
            host=config['postgresql']['HOST'],
            port=config['postgresql']['PORT'],
            user=config['postgresql']['USER_NAME'],
            password=config['postgresql']['PASSWD'],
            database='QA_knowledge'  # 连接到QA_knowledge数据库
        )
        print(f"✅ 成功连接到QA_knowledge数据库")
        return conn
    except Exception as e:
        print(f"❌ 连接QA_knowledge数据库失败: {e}")
        raise


def create_table_images_table(conn):
    """创建table_images表（支持多张图片）"""
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS table_images (
        table_name VARCHAR(255) PRIMARY KEY,
        table_title VARCHAR(500),
        table_description TEXT,
        images JSONB,
        reservoir_name VARCHAR(100),
        category_id VARCHAR(50),
        created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
    );
    
    CREATE INDEX IF NOT EXISTS idx_table_images_reservoir ON table_images (reservoir_name);
    CREATE INDEX IF NOT EXISTS idx_table_images_category ON table_images (category_id);
    """
    
    try:
        with conn.cursor() as cursor:
            cursor.execute(create_table_sql)
            conn.commit()
            print("✅ table_images表创建/检查成功（支持多图片）")
    except Exception as e:
        print(f"❌ 创建表失败: {e}")
        raise


def load_table_data(database: str) -> List[Dict]:
    """从JSON文件加载表格数据"""
    json_path = os.path.join(DATASET_DIR, database, f"{database}_tables.json")
    
    if not os.path.exists(json_path):
        print(f"⚠️ 找不到文件: {json_path}")
        return []
    
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    print(f"✅ 从{json_path}加载了{len(data)}个表格")
    return data


def find_table_images(table_title: str, database: str) -> List[str]:
    """查找表格对应的图片文件"""
    image_dir = os.path.join(DATASET_DIR, database, "tables_image")
    
    if not os.path.exists(image_dir):
        return []
    
    # 清理表格标题中的特殊字符
    clean_title = table_title.replace("/", "").replace("\\", "")
    
    image_paths = []
    for file in os.listdir(image_dir):
        if clean_title in file and file.lower().endswith((".png", ".jpg", ".jpeg")):
            full_path = os.path.join(image_dir, file)
            image_paths.append(full_path)
    
    return image_paths


def image_to_base64(image_path: str) -> str:
    """将图片转换为base64编码"""
    with open(image_path, "rb") as f:
        img_bytes = f.read()
        img_base64 = base64.b64encode(img_bytes).decode("utf-8")
    return f"data:image/jpeg;base64,{img_base64}"


def insert_table_images(conn, database: str):
    """插入表格图片数据（支持一个表格多张图片）"""
    # 加载表格数据
    tables_data = load_table_data(database)
    
    if not tables_data:
        print("⚠️ 没有表格数据可导入")
        return
    
    # 准备插入数据，使用字典合并同名表格
    # 格式: {table_name: {"title": ..., "description": ..., "images": [...]}}
    table_groups = {}
    
    for table in tables_data:
        table_title = table.get('table_title', '')
        table_name = table_title  # 使用标题作为表名
        
        # 如果这个表名已经处理过，跳过（避免重复）
        if table_name in table_groups:
            print(f"⏭️  跳过重复表格: {table_title} (已处理)")
            continue
        
        # 查找图片
        image_paths = find_table_images(table_title, database)
        
        if not image_paths:
            print(f"⚠️ 未找到表格 '{table_title}' 的图片")
            continue
        
        # 创建新记录
        table_groups[table_name] = {
            "title": table_title,
            "description": table.get('table_description', ''),
            "images": []
        }
        print(f"✅ 发现表格: {table_title}")
        
        # 将所有找到的图片都添加进去
        for idx, image_path in enumerate(image_paths, start=1):
            image_base64 = image_to_base64(image_path)
            
            # 判断是否是续表
            image_filename = os.path.basename(image_path)
            is_continued = any(x in image_filename for x in ['续表', '2', '-2', '_2'])
            
            image_info = {
                "url": image_base64,
                "index": len(table_groups[table_name]["images"]) + 1,
                "description": "续表" if is_continued else "主表",
                "filename": image_filename
            }
            
            table_groups[table_name]["images"].append(image_info)
            print(f"   📷 添加图片 {len(table_groups[table_name]['images'])}: {image_filename}")
    
    # 逐条插入
    if table_groups:
        insert_sql = """
        INSERT INTO table_images (
            table_name, table_title, table_description, 
            images, reservoir_name, category_id
        ) VALUES (%s, %s, %s, %s, %s, %s)
        ON CONFLICT (table_name) DO UPDATE SET
            table_title = EXCLUDED.table_title,
            table_description = EXCLUDED.table_description,
            images = EXCLUDED.images,
            reservoir_name = EXCLUDED.reservoir_name,
            category_id = EXCLUDED.category_id,
            updated_at = CURRENT_TIMESTAMP;
        """
        
        success_count = 0
        try:
            with conn.cursor() as cursor:
                for table_name, table_data in table_groups.items():
                    try:
                        # 将images列表转换为JSON字符串
                        images_json = json.dumps(table_data["images"], ensure_ascii=False)
                        
                        cursor.execute(insert_sql, (
                            table_name,
                            table_data["title"],
                            table_data["description"],
                            images_json,
                            database.replace('_', ' '),
                            'DHF'
                        ))
                        success_count += 1
                        print(f"✅ 成功插入: {table_name} (包含 {len(table_data['images'])} 张图片)")
                    except Exception as e:
                        print(f"❌ 插入 '{table_name}' 失败: {e}")
                        continue
                
                conn.commit()
                print(f"\n✅ 总计成功插入/更新 {success_count}/{len(table_groups)} 条表格数据")
        except Exception as e:
            print(f"❌ 插入数据失败: {e}")
            conn.rollback()
            raise


def main():
    """主函数"""
    print("=" * 60)
    print("表格图片数据导入工具")
    print("=" * 60)
    
    try:
        # 连接数据库
        conn = get_qa_database_connection()
        
        # 创建表
        create_table_images_table(conn)
        
        # 导入数据
        for database in DATABASES:
            print(f"\n📊 处理数据库: {database}")
            insert_table_images(conn, database)
        
        print("\n" + "=" * 60)
        print("✅ 所有表格图片数据导入完成！")
        print("=" * 60)
        
        # 查询统计
        with conn.cursor() as cursor:
            cursor.execute("SELECT COUNT(*) FROM table_images")
            count = cursor.fetchone()[0]
            print(f"\n📈 当前table_images表中共有 {count} 条记录")
        
        conn.close()
        
    except Exception as e:
        print(f"\n❌ 导入失败: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()

